import torch.nn as nn


class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # n*1*64*64
        self.c1 = nn.Sequential(
            nn.Conv2d(1, 12, kernel_size=(5, 5)),
            nn.ELU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=2))
        # n*12*30*30
        self.c2 = nn.Sequential(
            nn.Conv2d(12, 18, kernel_size=(5, 5)),
            nn.ELU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=2))
        # n*18*13*13
        self.c3 = nn.Sequential(
            nn.Conv2d(18, 27, kernel_size=(5, 5), stride=2),
            nn.ELU())
        # n*27*5*5
        self.fc4 = nn.Sequential(
            nn.Linear(27 * 5 * 5, 160),
            nn.Softsign())
        # n*160
        self.fc5 = nn.Sequential(
            nn.Linear(160, 40))
        # n*40

    def forward(self, x):
        x = self.c1(x)
        x = self.c2(x)
        x = self.c3(x)
        x = x.view(-1, 27 * 5 * 5)
        x = self.fc4(x)
        x = self.fc5(x)
        return x
